package net.minecraft.entity.ai.brain.task;

import java.util.Optional;
import net.minecraft.entity.EntityStatuses;
import net.minecraft.entity.ai.brain.MemoryModuleType;
import net.minecraft.entity.passive.VillagerEntity;
import net.minecraft.registry.Registries;
import net.minecraft.server.MinecraftServer;
import net.minecraft.util.math.GlobalPos;
import net.minecraft.village.VillagerProfession;

public class UpdateJobSiteTask {
	public static Task<VillagerEntity> create() {
		return TaskTriggerer.task(
			context -> context.group(context.queryMemoryValue(MemoryModuleType.POTENTIAL_JOB_SITE), context.queryMemoryOptional(MemoryModuleType.JOB_SITE))
				.apply(
					context,
					(potentialJobSite, jobSite) -> (world, entity, time) -> {
						GlobalPos globalPos = context.getValue(potentialJobSite);
						if (!globalPos.pos().isWithinDistance(entity.getPos(), 2.0) && !entity.isNatural()) {
							return false;
						} else {
							potentialJobSite.forget();
							jobSite.remember(globalPos);
							world.sendEntityStatus(entity, EntityStatuses.ADD_VILLAGER_HAPPY_PARTICLES);
							if (!entity.getVillagerData().profession().matchesKey(VillagerProfession.NONE)) {
								return true;
							} else {
								MinecraftServer minecraftServer = world.getServer();
								Optional.ofNullable(minecraftServer.getWorld(globalPos.dimension()))
									.flatMap(jobSiteWorld -> jobSiteWorld.getPointOfInterestStorage().getType(globalPos.pos()))
									.flatMap(
										poiType -> Registries.VILLAGER_PROFESSION
											.streamEntries()
											.filter(profession -> ((VillagerProfession)profession.value()).heldWorkstation().test(poiType))
											.findFirst()
									)
									.ifPresent(profession -> {
										entity.setVillagerData(entity.getVillagerData().withProfession(profession));
										entity.reinitializeBrain(world);
									});
								return true;
							}
						}
					}
				)
		);
	}
}
